from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
import numpy as np
from scipy.linalg import qr, svd, inv
import logging
from sklearn.neighbors import NearestNeighbors
import copy
import svgwrite
import os

def canoncorr(X0: np.array, Y0: np.array, fullReturn: bool = False) -> np.array:
    """
    Canonical Correlation Analysis (CCA) with added diagnostics and preprocessing
    
    Parameters:
    X, Y: (samples/observations) x (features) matrix
    fullReturn: whether all outputs should be returned or just `r` be returned
    
    Returns: A, B, r, U, V (if fullReturn is True) or just r (if fullReturn is False)
    A, B: Canonical coefficients for X and Y
    U, V: Canonical scores for the variables X and Y
    r: Canonical correlations
    """
    n, p1 = X0.shape
    p2 = Y0.shape[1]
    
    # Data diagnostics
    print(f"X shape: {X0.shape}, Y shape: {Y0.shape}")
    print(f"X condition number: {np.linalg.cond(X0)}")
    print(f"Y condition number: {np.linalg.cond(Y0)}")
    
    if p1 >= n or p2 >= n:
        logging.warning('Not enough samples, might cause problems')

    # Preprocessing: Standardize the variables
    X = (X0 - np.mean(X0, 0)) / np.std(X0, 0)
    Y = (Y0 - np.mean(Y0, 0)) / np.std(Y0, 0)
        
    # Factor the inputs, and find a full rank set of columns if necessary
    Q1, T11, perm1 = qr(X, mode='economic', pivoting=True)
    Q2, T22, perm2 = qr(Y, mode='economic', pivoting=True)

    # Determine ranks with a more stringent threshold
    tol = np.finfo(float).eps * 100  # Increased tolerance
    rankX = np.sum(np.abs(np.diagonal(T11)) > tol * np.abs(T11[0, 0]))
    rankY = np.sum(np.abs(np.diagonal(T22)) > tol * np.abs(T22[0, 0]))

    print(f"Rank of X: {rankX}, Rank of Y: {rankY}")

    if rankX == 0:
        raise ValueError('X has zero rank')
    elif rankX < p1:
        logging.warning('X is not full rank')
        Q1 = Q1[:, :rankX]
        T11 = T11[:rankX, :rankX]

    if rankY == 0:
        raise ValueError('Y has zero rank')
    elif rankY < p2:
        logging.warning('Y is not full rank')
        Q2 = Q2[:, :rankY]
        T22 = T22[:rankY, :rankY]

    # Compute canonical coefficients and canonical correlations
    d = min(rankX, rankY)
    L,D,M = svd(Q1.T @ Q2, full_matrices=True, check_finite=True, lapack_driver='gesdd')
    M = M.T

    A = inv(T11) @ L[:, :d] * np.sqrt(n - 1)
    B = inv(T22) @ M[:, :d] * np.sqrt(n - 1)
    r = D[:d]
    
    # Remove roundoff errors
    r = np.clip(r, 0, 1)

    if not fullReturn:
        return r

    # Put coefficients back to their full size and correct order
    A_full = np.zeros((p1, d))
    A_full[perm1, :] = np.vstack((A, np.zeros((p1 - rankX, d))))
    
    B_full = np.zeros((p2, d))
    B_full[perm2, :] = np.vstack((B, np.zeros((p2 - rankY, d))))

    # Compute the canonical variates
    U = X @ A_full
    V = Y @ B_full

    return A_full, B_full, r, U, V


def plot_lollipop(scores, title="Lollipop Plot of Scores", xlabel="Index", ylabel="Score", figsize=(12, 6), y_lim=(0, 1.1)):
    """
    绘制带连线和数值标注的棒棒糖图，数值保留两位小数，x轴标签从1开始
    """
    
    # 创建对应的 x 值
    x = np.arange(len(scores))

    # 创建图表
    fig, ax = plt.subplots(figsize=figsize)

    # 绘制垂直线
    ax.vlines(x, 0, scores, colors='gray', lw=1, alpha=0.5)

    # 绘制数据点并连线
    ax.plot(x, scores, color='black', marker='o', markersize=8, linestyle='-', linewidth=1)

    # 添加数值标注，保留两位小数
    for i, score in enumerate(scores):
        ax.annotate(f'{score:.2f}', (i, score), textcoords="offset points", 
                    xytext=(0,10), ha='center', va='bottom', fontsize=8)

    # 设置坐标轴范围
    ax.set_xlim(-0.5, len(scores) - 0.5)
    ax.set_ylim(0, max(scores) * 1.2)  # 增加上界以容纳标注

    # 设置标题和标签
    ax.set_title(title)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)

    # 移除顶部和右侧边框
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)

    # 移除网格线
    ax.grid(False)

    # 调整 x 轴刻度，标签从1开始
    ax.set_xticks(x)
    ax.set_xticklabels(range(1, len(scores) + 1))

    # 设置 y 范围 0-1
    ax.set_ylim(y_lim)

    plt.tight_layout()
    plt.show()


def plot_comparison_lollipop(scores1, scores2, label1="Series 1", label2="Series 2", title="Comparison Lollipop Plot", xlabel="Index", ylabel="Score", figsize=(12, 6), y_lim=(0, 1.1)):
    """
    绘制两组数据对比的棒棒糖图
    """
    # 确保两组数据等长
    min_len = min(len(scores1), len(scores2))
    scores1 = scores1[:min_len]
    scores2 = scores2[:min_len]
    
    # 创建对应的 x 值
    x = np.arange(len(scores1))

    # 创建图表
    fig, ax = plt.subplots(figsize=figsize)

    # 绘制第一组数据
    ax.vlines(x, 0, scores1, colors='blue', lw=1, alpha=0.5)
    ax.plot(x, scores1, color='blue', marker='o', markersize=8, linestyle='-', linewidth=1, label=label1)
    for i, score in enumerate(scores1):
        ax.annotate(f'{score:.2f}', (i, score), textcoords="offset points", 
                    xytext=(0,10), ha='center', va='bottom', fontsize=8, color='blue')

    # 绘制第二组数据
    ax.vlines(x, 0, scores2, colors='red', lw=1, alpha=0.5)
    ax.plot(x, scores2, color='red', marker='o', markersize=8, linestyle='-', linewidth=1, label=label2)
    for i, score in enumerate(scores2):
        ax.annotate(f'{score:.2f}', (i, score), textcoords="offset points", 
                    xytext=(0,-15), ha='center', va='top', fontsize=8, color='red')

    # 设置坐标轴范围
    ax.set_xlim(-0.5, len(scores1) - 0.5)
    ax.set_ylim(y_lim)

    # 设置标题和标签
    ax.set_title(title)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)

    # 移除顶部和右侧边框
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)

    # 移除网格线
    ax.grid(False)

    # 调整 x 轴刻度，标签从1开始
    ax.set_xticks(x)
    ax.set_xticklabels(range(1, len(scores1) + 1))

    # 添加图例
    ax.legend()

    plt.tight_layout()
    plt.show()


def save_comparison_lollipop_svg(scores1, scores2, label1="Series 1", label2="Series 2", title="Comparison Lollipop Plot", xlabel="Index", ylabel="Score", nn_type="vanilla"):
    """
    将CCA对比棒棒糖图保存为高质量SVG格式，采用科学论文风格
    """
    try:
        print("Starting SVG save process...")
        
        # 确保两组数据等长
        min_len = min(len(scores1), len(scores2))
        scores1 = scores1[:min_len]
        scores2 = scores2[:min_len]
        
        # SVG画布尺寸 - 适合科学论文的比例
        svg_width = 800
        svg_height = 500
        margin_left = 80
        margin_right = 40
        margin_top = 60
        margin_bottom = 80
        plot_width = svg_width - margin_left - margin_right
        plot_height = svg_height - margin_top - margin_bottom
        
        # 创建SVG画布
        svg_path = f"./logs/cca_comparison_{nn_type}.svg"
        dwg = svgwrite.Drawing(svg_path, size=(svg_width, svg_height))
        
        # 背景 - 纯白色，符合科学论文标准
        dwg.add(dwg.rect(insert=(0, 0), size=(svg_width, svg_height), fill='white'))
        
        # 数据范围计算
        max_score = max(max(scores1), max(scores2))
        y_upper = 1.1  # 固定上限为1.1，适合相关系数
        
        # 坐标转换函数
        def data_to_svg_coords(x_idx, y_val):
            if len(scores1) <= 1:
                svg_x = margin_left + plot_width // 2
            else:
                svg_x = margin_left + (x_idx / (len(scores1) - 1)) * plot_width
            svg_y = margin_top + plot_height - (y_val / y_upper) * plot_height
            return svg_x, svg_y
        
        # 绘制坐标轴边框 - 深灰色，更专业
        axis_color = 'rgb(60,60,60)'
        dwg.add(dwg.rect(insert=(margin_left, margin_top), size=(plot_width, plot_height),
                        fill='none', stroke=axis_color, stroke_width=2))
        
        # 科学论文级别的颜色方案
        # 第一组数据：深蓝色（表示训练模型，更稳定）
        color1_main = 'rgb(31,119,180)'      # 深蓝色主色
        color1_light = 'rgb(174,199,232)'    # 浅蓝色辅助
        color1_dark = 'rgb(23,90,140)'       # 深蓝色强调
        
        # 第二组数据：深橙色（表示对照组，需要对比）
        color2_main = 'rgb(255,127,14)'      # 深橙色主色  
        color2_light = 'rgb(255,187,120)'    # 浅橙色辅助
        color2_dark = 'rgb(196,95,11)'       # 深橙色强调
        
        # 绘制第一组数据的连接线（柔和的渐变效果）
        for i in range(len(scores1)):
            svg_x, svg_y_bottom = data_to_svg_coords(i, 0)
            svg_x, svg_y_top = data_to_svg_coords(i, scores1[i])
            
            # 垂直连接线 - 使用渐变效果
            dwg.add(dwg.line(start=(svg_x, svg_y_bottom), end=(svg_x, svg_y_top),
                            stroke=color1_light, stroke_width=3, opacity=0.8))
        
        # 绘制第二组数据的连接线
        for i in range(len(scores2)):
            svg_x, svg_y_bottom = data_to_svg_coords(i, 0)
            svg_x, svg_y_top = data_to_svg_coords(i, scores2[i])
            
            # 垂直连接线
            dwg.add(dwg.line(start=(svg_x, svg_y_bottom), end=(svg_x, svg_y_top),
                            stroke=color2_light, stroke_width=3, opacity=0.8))
        
        # 绘制第一组数据的主线条（连接各点）
        if len(scores1) > 1:
            points1 = []
            for i in range(len(scores1)):
                svg_x, svg_y = data_to_svg_coords(i, scores1[i])
                points1.append((svg_x, svg_y))
            
            # 主线条 - 更粗，体现数据连续性
            dwg.add(dwg.polyline(points=points1, fill='none', 
                               stroke=color1_main, stroke_width=3, opacity=0.9))
        
        # 绘制第二组数据的主线条
        if len(scores2) > 1:
            points2 = []
            for i in range(len(scores2)):
                svg_x, svg_y = data_to_svg_coords(i, scores2[i])
                points2.append((svg_x, svg_y))
            
            dwg.add(dwg.polyline(points=points2, fill='none', 
                               stroke=color2_main, stroke_width=3, opacity=0.9))
        
        # 绘制数据点 - 加上阴影效果和边框
        for i, score in enumerate(scores1):
            svg_x, svg_y = data_to_svg_coords(i, score)
            
            # 阴影效果
            dwg.add(dwg.circle(center=(svg_x + 2, svg_y + 2), r=6, 
                             fill='rgb(200,200,200)', opacity=0.3))
            # 外圈（深色边框）
            dwg.add(dwg.circle(center=(svg_x, svg_y), r=6, 
                             fill=color1_dark, stroke='none'))
            # 内圈（主色）
            dwg.add(dwg.circle(center=(svg_x, svg_y), r=4, 
                             fill=color1_main, stroke='none'))
            # 高光点
            dwg.add(dwg.circle(center=(svg_x - 1, svg_y - 1), r=1.5, 
                             fill='rgb(255,255,255)', opacity=0.8))
            
            # 数值标注 - 科学论文风格
            dwg.add(dwg.text(f'{score:.2f}', insert=(svg_x, svg_y - 12), 
                            text_anchor='middle', font_size=10, font_family='Arial, sans-serif',
                            fill=color1_dark, font_weight='bold'))
        
        # 绘制第二组数据点
        for i, score in enumerate(scores2):
            svg_x, svg_y = data_to_svg_coords(i, score)
            
            # 阴影效果
            dwg.add(dwg.circle(center=(svg_x + 2, svg_y + 2), r=6, 
                             fill='rgb(200,200,200)', opacity=0.3))
            # 外圈（深色边框）
            dwg.add(dwg.circle(center=(svg_x, svg_y), r=6, 
                             fill=color2_dark, stroke='none'))
            # 内圈（主色）
            dwg.add(dwg.circle(center=(svg_x, svg_y), r=4, 
                             fill=color2_main, stroke='none'))
            # 高光点
            dwg.add(dwg.circle(center=(svg_x - 1, svg_y - 1), r=1.5, 
                             fill='rgb(255,255,255)', opacity=0.8))
            
            # 数值标注 - 位置稍微下移避免重叠
            dwg.add(dwg.text(f'{score:.2f}', insert=(svg_x, svg_y + 20), 
                            text_anchor='middle', font_size=10, font_family='Arial, sans-serif',
                            fill=color2_dark, font_weight='bold'))
        
        # 添加坐标轴标签 - 科学论文字体
        label_color = 'rgb(40,40,40)'
        dwg.add(dwg.text(xlabel, insert=(margin_left + plot_width//2, svg_height - 20), 
                        text_anchor='middle', font_size=14, font_family='Arial, sans-serif',
                        fill=label_color, font_weight='normal'))
        
        # Y轴标签（旋转90度）
        dwg.add(dwg.text(ylabel, insert=(20, margin_top + plot_height//2), 
                        text_anchor='middle', font_size=14, font_family='Arial, sans-serif',
                        fill=label_color, font_weight='normal',
                        transform=f'rotate(-90 20 {margin_top + plot_height//2})'))
        
        # 添加标题 - 更大字体，专业风格
        dwg.add(dwg.text(title, insert=(svg_width//2, 30), 
                        text_anchor='middle', font_size=16, font_family='Arial, sans-serif',
                        fill=label_color, font_weight='bold'))
        
        # 添加图例 - 科学论文风格
        legend_x = margin_left + plot_width - 200
        legend_y = margin_top + 50
        
        # 图例背景框
        dwg.add(dwg.rect(insert=(legend_x - 10, legend_y - 10), size=(190, 60),
                        fill='rgb(250,250,250)', stroke='rgb(200,200,200)', 
                        stroke_width=1, opacity=0.95))
        
        # 第一组图例
        dwg.add(dwg.circle(center=(legend_x + 10, legend_y + 10), r=5, fill=color1_main))
        dwg.add(dwg.text(label1, insert=(legend_x + 25, legend_y + 15), 
                        font_size=12, font_family='Arial, sans-serif',
                        fill=label_color, font_weight='normal'))
        
        # 第二组图例
        dwg.add(dwg.circle(center=(legend_x + 10, legend_y + 30), r=5, fill=color2_main))
        dwg.add(dwg.text(label2, insert=(legend_x + 25, legend_y + 35), 
                        font_size=12, font_family='Arial, sans-serif',
                        fill=label_color, font_weight='normal'))
        
        # 添加x轴刻度标签
        for i in range(len(scores1)):
            svg_x, _ = data_to_svg_coords(i, 0)
            dwg.add(dwg.text(str(i + 1), insert=(svg_x, margin_top + plot_height + 15), 
                            text_anchor='middle', font_size=10, font_family='Arial, sans-serif',
                            fill=label_color))
        
        # 保存SVG文件
        dwg.save()
        print(f"CCA comparison SVG saved to {svg_path}")
        
    except Exception as e:
        print(f"Error saving CCA comparison SVG: {e}")
        import traceback
        traceback.print_exc()


def plot_2d_cca_separate(U_grouped_averaged, V_grouped_averaged, colors, nn_type):
    """
    绘制2D版本的CCA点云图，分别显示Neural Dynamics和Behavior
    """
    # 路径长度对应关系 (索引0-9对应路径长度5-14)
    path_lengths = [5, 6, 7, 8, 9, 10, 11, 12, 13, 14]
    
    # CM0/1 组合
    # 绘制Neural Dynamics (CM0/1)
    fig1, ax1 = plt.subplots(figsize=(10, 8))
    
    for j in range(10):
        color_rgb = colors[j]
        ax1.scatter(U_grouped_averaged[j][:, 0], U_grouped_averaged[j][:, 1], 
                   color=color_rgb, marker='o', s=30, alpha=0.7, label=f'Length {path_lengths[j]}')
    
    ax1.set_title('Limit Cycle Distribution (CM0/1)', fontsize=14, fontweight='bold')
    ax1.set_xlabel('Canonical mode 0', fontsize=12)
    ax1.set_ylabel('Canonical mode 1', fontsize=12)
    ax1.grid(True, alpha=0.3)
    ax1.set_aspect('equal', adjustable='box')
    
    # 添加图例到右上角
    ax1.legend(bbox_to_anchor=(1.05, 1), loc='upper left', title='Path Length')
    
    plt.tight_layout()
    plt.show()
    
    # 绘制Behavior Structure (CM0/1)
    fig2, ax2 = plt.subplots(figsize=(10, 8))
    
    for j in range(10):
        color_rgb = colors[j]
        ax2.scatter(V_grouped_averaged[j][:, 0], V_grouped_averaged[j][:, 1], 
                   color=color_rgb, marker='o', s=30, alpha=0.7, label=f'Length {path_lengths[j]}')
    
    ax2.set_title('Behavior Structure (CM0/1)', fontsize=14, fontweight='bold')
    ax2.set_xlabel('Canonical mode 0', fontsize=12)
    ax2.set_ylabel('Canonical mode 1', fontsize=12)
    ax2.grid(True, alpha=0.3)
    ax2.set_aspect('equal', adjustable='box')
    
    # 添加图例到右上角
    ax2.legend(bbox_to_anchor=(1.05, 1), loc='upper left', title='Path Length')
    
    plt.tight_layout()
    plt.show()
    
    # CM3/4 组合
    # 绘制Neural Dynamics (CM3/4)
    fig3, ax3 = plt.subplots(figsize=(10, 8))
    
    for j in range(10):
        color_rgb = colors[j]
        ax3.scatter(U_grouped_averaged[j][:, 3], U_grouped_averaged[j][:, 4], 
                   color=color_rgb, marker='o', s=30, alpha=0.7, label=f'Length {path_lengths[j]}')
    
    ax3.set_title('Limit Cycle Distribution (CM3/4)', fontsize=14, fontweight='bold')
    ax3.set_xlabel('Canonical mode 3', fontsize=12)
    ax3.set_ylabel('Canonical mode 4', fontsize=12)
    ax3.grid(True, alpha=0.3)
    ax3.set_aspect('equal', adjustable='box')
    
    # 添加图例到右上角
    ax3.legend(bbox_to_anchor=(1.05, 1), loc='upper left', title='Path Length')
    
    plt.tight_layout()
    plt.show()
    
    # 绘制Behavior Structure (CM3/4)
    fig4, ax4 = plt.subplots(figsize=(10, 8))
    
    for j in range(10):
        color_rgb = colors[j]
        ax4.scatter(V_grouped_averaged[j][:, 3], V_grouped_averaged[j][:, 4], 
                   color=color_rgb, marker='o', s=30, alpha=0.7, label=f'Length {path_lengths[j]}')
    
    ax4.set_title('Behavior Structure (CM3/4)', fontsize=14, fontweight='bold')
    ax4.set_xlabel('Canonical mode 3', fontsize=12)
    ax4.set_ylabel('Canonical mode 4', fontsize=12)
    ax4.grid(True, alpha=0.3)
    ax4.set_aspect('equal', adjustable='box')
    
    # 添加图例到右上角
    ax4.legend(bbox_to_anchor=(1.05, 1), loc='upper left', title='Path Length')
    
    plt.tight_layout()
    plt.show()
    
    # CM3/5 组合
    # 绘制Neural Dynamics (CM3/5)
    fig5, ax5 = plt.subplots(figsize=(10, 8))
    
    for j in range(10):
        color_rgb = colors[j]
        ax5.scatter(U_grouped_averaged[j][:, 3], U_grouped_averaged[j][:, 5], 
                   color=color_rgb, marker='o', s=30, alpha=0.7, label=f'Length {path_lengths[j]}')
    
    ax5.set_title('Limit Cycle Distribution (CM3/5)', fontsize=14, fontweight='bold')
    ax5.set_xlabel('Canonical mode 3', fontsize=12)
    ax5.set_ylabel('Canonical mode 5', fontsize=12)
    ax5.grid(True, alpha=0.3)
    ax5.set_aspect('equal', adjustable='box')
    
    # 添加图例到右上角
    ax5.legend(bbox_to_anchor=(1.05, 1), loc='upper left', title='Path Length')
    
    plt.tight_layout()
    plt.show()
    
    # 绘制Behavior Structure (CM3/5)
    fig6, ax6 = plt.subplots(figsize=(10, 8))
    
    for j in range(10):
        color_rgb = colors[j]
        ax6.scatter(V_grouped_averaged[j][:, 3], V_grouped_averaged[j][:, 5], 
                   color=color_rgb, marker='o', s=30, alpha=0.7, label=f'Length {path_lengths[j]}')
    
    ax6.set_title('Behavior Structure (CM3/5)', fontsize=14, fontweight='bold')
    ax6.set_xlabel('Canonical mode 3', fontsize=12)
    ax6.set_ylabel('Canonical mode 5', fontsize=12)
    ax6.grid(True, alpha=0.3)
    ax6.set_aspect('equal', adjustable='box')
    
    # 添加图例到右上角
    ax6.legend(bbox_to_anchor=(1.05, 1), loc='upper left', title='Path Length')
    
    plt.tight_layout()
    plt.show()


def save_single_cca_svg(data_grouped_averaged, colors, nn_type, data_type, cm1, cm2, title_suffix):
    """
    保存单个CCA点云图为SVG格式
    """
    # 路径长度对应关系 (索引0-9对应路径长度5-14)
    path_lengths = [5, 6, 7, 8, 9, 10, 11, 12, 13, 14]
    
    # SVG画布尺寸
    svg_width = 700
    svg_height = 600
    subplot_width = 500
    subplot_height = 500
    margin = 50
    
    # 创建SVG文件路径
    svg_path = f"./logs/2d_cca_{data_type}_cm{cm1}{cm2}_{nn_type}.svg"
    dwg = svgwrite.Drawing(svg_path, size=(svg_width, svg_height))
    
    # 背景
    dwg.add(dwg.rect(insert=(0, 0), size=(svg_width, svg_height), fill='white'))
    
    # 计算数据范围
    data_min_x = np.min(data_grouped_averaged[:, :, cm1])
    data_max_x = np.max(data_grouped_averaged[:, :, cm1])
    data_min_y = np.min(data_grouped_averaged[:, :, cm2])
    data_max_y = np.max(data_grouped_averaged[:, :, cm2])
    
    data_range_x = data_max_x - data_min_x
    data_range_y = data_max_y - data_min_y
    data_padding_x = data_range_x * 0.1
    data_padding_y = data_range_y * 0.1
    
    def data_to_svg_coords(x, y):
        svg_x = margin + (x - (data_min_x - data_padding_x)) / (data_range_x + 2*data_padding_x) * subplot_width
        svg_y = margin + (1 - (y - (data_min_y - data_padding_y)) / (data_range_y + 2*data_padding_y)) * subplot_height
        return svg_x, svg_y
    
    # 绘制边框
    dwg.add(dwg.rect(insert=(margin, margin), size=(subplot_width, subplot_height), 
                   fill='none', stroke='black', stroke_width=2))
    
    # 绘制数据点
    for j in range(10):
        color_rgb = colors[j]
        color_255 = [int(c * 255) for c in color_rgb]
        svg_color = f'rgb({color_255[0]},{color_255[1]},{color_255[2]})'
        
        for point in data_grouped_averaged[j]:
            svg_x, svg_y = data_to_svg_coords(point[cm1], point[cm2])
            dwg.add(dwg.circle(center=(svg_x, svg_y), r=3, 
                             fill=svg_color, opacity=0.7))
    
    # 添加标题
    title_text = f"{title_suffix} (CM{cm1}/{cm2})"
    dwg.add(dwg.text(title_text, 
                   insert=(margin + subplot_width//2, margin - 10), text_anchor='middle', 
                   font_size=16, font_weight='bold', fill='black'))
    
    # 添加坐标轴标签
    dwg.add(dwg.text(f'Canonical mode {cm1}', 
                   insert=(margin + subplot_width//2, margin + subplot_height + 40), text_anchor='middle', 
                   font_size=12, fill='black'))
    
    dwg.add(dwg.text(f'Canonical mode {cm2}', 
                   insert=(15, margin + subplot_height//2), text_anchor='middle', 
                   font_size=12, fill='black', transform=f'rotate(-90 15 {margin + subplot_height//2})'))
    
    # 添加图例
    legend_start_x = margin + subplot_width + 20
    legend_start_y = margin + 20
    for j in range(10):
        color_rgb = colors[j]
        color_255 = [int(c * 255) for c in color_rgb]
        svg_color = f'rgb({color_255[0]},{color_255[1]},{color_255[2]})'
        
        legend_y = legend_start_y + j * 20
        dwg.add(dwg.circle(center=(legend_start_x + 10, legend_y), r=6, 
                         fill=svg_color, opacity=0.7))
        dwg.add(dwg.text(f'Length {path_lengths[j]}', 
                       insert=(legend_start_x + 25, legend_y + 5), 
                       font_size=10, fill='black'))
    
    # 图例标题
    dwg.add(dwg.text('Path Length', 
                   insert=(legend_start_x + 10, legend_start_y - 10), 
                   font_size=12, font_weight='bold', fill='black'))
    
    dwg.save()
    print(f"{title_suffix} CM{cm1}/{cm2} SVG saved to {svg_path}")


def save_2d_cca_svg(U_grouped_averaged, V_grouped_averaged, colors, nn_type):
    """
    将2D版本的CCA点云图保存为多个SVG文件，包括CM0/1、CM3/4、CM3/5
    """
    # 保存所有组合
    canonical_mode_pairs = [(0, 1), (3, 4), (3, 5)]
    
    for cm1, cm2 in canonical_mode_pairs:
        # 保存Limit Cycle Distribution
        save_single_cca_svg(U_grouped_averaged, colors, nn_type, 
                          "limit_cycle_distribution", cm1, cm2, 
                          "Limit Cycle Distribution")
        
        # 保存Behavior Structure
        save_single_cca_svg(V_grouped_averaged, colors, nn_type, 
                          "behavior_structure", cm1, cm2, 
                          "Behavior Structure")


def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

def preprocess(trjs, max_length):
    processed_trjs = []
    for i in range(trjs.shape[0]):
        trj = trjs[i]
        if trj.shape[0] < max_length:
            last_element = trj[-1]
            processed_trj = np.concatenate([trj, np.repeat(last_element[np.newaxis,:], max_length - trj.shape[0], axis=0)], axis=0)
            processed_trjs.append(processed_trj)
        else:
            processed_trjs.append(trj)
    processed_trjs = np.array(processed_trjs)
    return processed_trjs
    
@jax.jit
def build_radiance_field(p):
    img = jnp.zeros((21, 21))
    top = 10
    bottom = 0
    effective_radius = 21*1.414
    px, py = p[0], p[1]

    x_coord_map = jnp.arange(img.shape[0])
    x_coord_map = jnp.repeat(x_coord_map[:, jnp.newaxis], img.shape[1], axis=1)

    y_coord_map = jnp.arange(img.shape[1])
    y_coord_map = jnp.repeat(y_coord_map[jnp.newaxis, :], img.shape[0], axis=0)
    
    dist = jnp.sqrt((x_coord_map - px)**2 + (y_coord_map - py)**2)
    img = jnp.where(dist < effective_radius,
                    (top - bottom) * (effective_radius - dist) / effective_radius + bottom,
                    img)
    return img

@jax.jit
def get_max_radiance_field(imgs):
    img = jnp.max(imgs, axis=0)
    return img

@jax.jit
def build_ridge(A):
    # 将 A 的第一个点对齐到 (10,10)
    A = A-A[0]+jnp.array([10,10])
    # 使用 jnp.map 来并行计算每个辐射场图像
    imgs = jax.vmap(build_radiance_field)(A)
    img = get_max_radiance_field(imgs)
    return img

build_ridge_vmap = jax.vmap(build_ridge)


def compute_trajectory_manifold():
    rpl_config = ReplayConfig()
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)

    args = parser.parse_args()

    nn_type = ''
    if args.nn_type == "vanilla":
        nn_type = "vanilla"
    elif args.nn_type == "gru":
        nn_type = "gru"

    def load_data_and_compute(nn_type, seq_len, redundancy, diverse_set_capacity):
        rnn_limit_rings_file_name = "./logs/rnn_limit_rings_of_best_estimation_" + nn_type + "_" + str(seq_len) + "_" + str(redundancy) + "_" + str(diverse_set_capacity) + ".npz"
        rnn_limit_rings_of_best_estimation_file = np.load(rnn_limit_rings_file_name)
        matrix_names = rnn_limit_rings_of_best_estimation_file.files
        rnn_limit_rings_of_best_estimation = []

        for name in matrix_names:
            matrix = rnn_limit_rings_of_best_estimation_file[name]
            rnn_limit_rings_of_best_estimation.append(matrix)

        rnn_limit_rings_of_best_estimation_center = []
        for i in range(len(rnn_limit_rings_of_best_estimation)):
            if rnn_limit_rings_of_best_estimation[i].shape[0] == 0:
                rnn_limit_rings_of_best_estimation_center.append(np.full((128,), 0))
            else:
                rnn_limit_rings_of_best_estimation_center.append(np.mean(rnn_limit_rings_of_best_estimation[i], axis=(0,1)))

        rnn_limit_rings_of_best_estimation_center = np.array(rnn_limit_rings_of_best_estimation_center)

        file_name = "obs_data_" + nn_type + "_" + str(seq_len) + "_" + str(redundancy) + "_" + str(diverse_set_capacity) + ".npz"
        obs_file = np.load("./logs/" + file_name)
        diverse_set_trajectoies = obs_file["diverse_set_trajectoies"]

        return rnn_limit_rings_of_best_estimation_center, diverse_set_trajectoies
    
    configs = [
        [nn_type, 5, 1, 400],
        [nn_type, 6, 1, 400],
        [nn_type, 7, 1, 400],
        [nn_type, 8, 1, 400],
        [nn_type, 9, 1, 400],
        [nn_type, 10, 1, 400],
        [nn_type, 11, 1, 400],
        [nn_type, 12, 1, 400],
        [nn_type, 13, 1, 400],
        [nn_type, 14, 1, 400],
    ]

    diverse_set_trajectoies = []
    for i in range(len(configs)):
        _, dts = load_data_and_compute(configs[i][0], configs[i][1], configs[i][2], configs[i][3])
        diverse_set_trajectoies.append(dts)

    diverse_set_trajectoies_mat = []
    for i in range(len(diverse_set_trajectoies)):
        length_aligned_trj = preprocess(diverse_set_trajectoies[i], configs[-1][1])
        diverse_set_trajectoies_mat.append(length_aligned_trj)
    diverse_set_trajectoies_mat = np.concatenate(diverse_set_trajectoies_mat, axis=0)

    ridge_images = build_ridge_vmap(diverse_set_trajectoies_mat)
    ridge_images = ridge_images.reshape(ridge_images.shape[0], ridge_images.shape[1] * ridge_images.shape[2])

    return copy.deepcopy(ridge_images)
    

def compute_policyring_manifold():
    rpl_config = ReplayConfig()
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)

    args = parser.parse_args()

    nn_type = ''
    if args.nn_type == "vanilla":
        nn_type = "vanilla"
    elif args.nn_type == "gru":
        nn_type = "gru"

    def load_data_and_compute(nn_type, seq_len, redundancy, diverse_set_capacity):
        rnn_limit_rings_file_name = "./logs/rnn_limit_rings_of_best_estimation_" + nn_type + "_" + str(seq_len) + "_" + str(redundancy) + "_" + str(diverse_set_capacity) + ".npz"
        rnn_limit_rings_of_best_estimation_file = np.load(rnn_limit_rings_file_name)
        matrix_names = rnn_limit_rings_of_best_estimation_file.files
        rnn_limit_rings_of_best_estimation = []

        for name in matrix_names:
            matrix = rnn_limit_rings_of_best_estimation_file[name]
            rnn_limit_rings_of_best_estimation.append(matrix)

        rnn_limit_rings_of_best_estimation_center = []
        for i in range(len(rnn_limit_rings_of_best_estimation)):
            if rnn_limit_rings_of_best_estimation[i].shape[0] == 0:
                rnn_limit_rings_of_best_estimation_center.append(np.full((128,), 0))
            else:
                rnn_limit_rings_of_best_estimation_center.append(np.mean(rnn_limit_rings_of_best_estimation[i], axis=(0,1)))
        rnn_limit_rings_of_best_estimation_center = np.array(rnn_limit_rings_of_best_estimation_center)

        return rnn_limit_rings_of_best_estimation_center
    
    configs = [
        [nn_type, 5, 1, 400],
        [nn_type, 6, 1, 400],
        [nn_type, 7, 1, 400],
        [nn_type, 8, 1, 400],
        [nn_type, 9, 1, 400],
        [nn_type, 10, 1, 400],
        [nn_type, 11, 1, 400],
        [nn_type, 12, 1, 400],
        [nn_type, 13, 1, 400],
        [nn_type, 14, 1, 400],
    ]

    rnn_limit_rings_of_best_estimation_centers = []
    ring_lengths = []
    for i in range(len(configs)):
        _rnn_limit_rings_of_best_estimation_centers = load_data_and_compute(configs[i][0], configs[i][1], configs[i][2], configs[i][3])
        rnn_limit_rings_of_best_estimation_centers.append(_rnn_limit_rings_of_best_estimation_centers)
        ring_length = configs[i][1]
        data_cap = _rnn_limit_rings_of_best_estimation_centers.shape[0]
        ring_length_ = np.array([ring_length for j in range(data_cap)])
        ring_lengths.append(ring_length_)

    ring_lengths = np.array(ring_lengths)
    ring_lengths = ring_lengths.reshape(ring_lengths.shape[0]*ring_lengths.shape[1])

    rnn_limit_rings_of_best_estimation_center_mat = np.concatenate(rnn_limit_rings_of_best_estimation_centers, axis=0)

    return copy.deepcopy(rnn_limit_rings_of_best_estimation_center_mat)


if __name__ == "__main__":
    rpl_config = ReplayConfig()
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)

    args = parser.parse_args()

    ridge_images = compute_trajectory_manifold()
    rnn_limit_rings_of_best_estimation_center_mat = compute_policyring_manifold()

    print("shape of ridge_images: ", ridge_images.shape)
    print("shape of rnn_limit_rings_of_best_estimation_center_mat: ", rnn_limit_rings_of_best_estimation_center_mat.shape)

    # 在 ridge_images 中加入平均噪声，分布在 [-0.001, 0.001] 区间
    noise = np.random.uniform(-0.001, 0.001, size=ridge_images.shape)
    ridge_images = ridge_images + noise
    # 在 rnn_limit_rings_of_best_estimation_center_mat 中加入平均噪声，分布在 [-0.001, 0.001] 区间
    noise = np.random.uniform(-0.001, 0.001, size=rnn_limit_rings_of_best_estimation_center_mat.shape)
    rnn_limit_rings_of_best_estimation_center_mat = rnn_limit_rings_of_best_estimation_center_mat + noise

    train_view_1 = rnn_limit_rings_of_best_estimation_center_mat
    train_view_2 = ridge_images

    print("shape of train_view_1: ", train_view_1.shape)
    print("shape of train_view_2: ", train_view_2.shape)

    # 对 train_view_1/2 进行 pca
    pca1 = PCA()
    pca1.fit(train_view_1)
    train_view_1 = pca1.transform(train_view_1)
    pca2 = PCA()
    pca2.fit(train_view_2)
    train_view_2 = pca2.transform(train_view_2)

    # using CCA from article 'Preserved neural dynamics across animals performing similar behaviour'
    A, B, r, U, V = canoncorr(train_view_1, train_view_2, fullReturn=True)
    print(A.shape, B.shape, r.shape, U.shape, V.shape)

    scores = r[:10]
    print("sum of scores: ", np.sum(scores))

    # 绘制 CCA 棒棒糖图
    plot_lollipop(scores, title="Canonical Correlation Analysis (CCA) of Ridge Representation and Policy Ring Manifolds "+' ('+args.nn_type+')', xlabel="Canonical mode", ylabel="Correlation", figsize=(12, 6))
    
    # 加载对比数据并绘制对比图
    try:
        control_scores_path = f"./logs/cca_scores_{args.nn_type}.npy"
        if os.path.exists(control_scores_path):
            control_scores = np.load(control_scores_path)
            print(f"Loaded control scores from {control_scores_path}")
            print(f"Control scores: {control_scores}")
            
            # 绘制对比棒棒糖图
            plot_comparison_lollipop(scores, control_scores, 
                                   label1="Policy Ring Manifolds", 
                                   label2="Action Constraint Manifolds", 
                                   title=f"CCA Comparison: Policy Ring vs Action Constraint Manifolds ({args.nn_type})", 
                                   xlabel="Canonical mode", 
                                   ylabel="Correlation", 
                                   figsize=(12, 6))
            
            # 保存对比图为SVG格式
            save_comparison_lollipop_svg(scores, control_scores, 
                                       label1="Policy Ring Manifolds", 
                                       label2="Action Constraint Manifolds", 
                                       title=f"CCA Comparison: Policy Ring vs Action Constraint Manifolds ({args.nn_type})", 
                                       xlabel="Canonical mode", 
                                       ylabel="Correlation", 
                                       nn_type=args.nn_type)
        else:
            print(f"Control scores file not found: {control_scores_path}")
    except Exception as e:
        print(f"Error loading control scores: {e}")

    # 仅保留高度相关的前10个模式
    U = U[:, :10]
    V = V[:, :10]
        
    # 对 U和V 进行分组
    U_grouped = []
    V_grouped = []
    for j in range(10):
        start = j * (U.shape[0]//10)
        end = (j+1) * (U.shape[0]//10)
        U_grouped.append(U[start:end])
        V_grouped.append(V[start:end])
    U_grouped = np.array(U_grouped)
    V_grouped = np.array(V_grouped)

    def neighborhood_average(grouped_data, n_neighbors=50):
        result = []
        for group in grouped_data:
            nn = NearestNeighbors(n_neighbors=n_neighbors+1, metric='euclidean')
            nn.fit(group)
            distances, indices = nn.kneighbors(group)
            
            averaged_group = np.array([
                np.mean(group[indices[i][1:]], axis=0)  # 排除自身，取后面50个邻居
                for i in range(len(group))
            ])
            
            result.append(averaged_group)
        
        return np.array(result)

    U_grouped_averaged = neighborhood_average(U_grouped)
    V_grouped_averaged = neighborhood_average(V_grouped)

    colors = []
    np.random.seed(100)
    for i in range(15):
        colors.append(np.random.rand(3))

    # 首先绘制2D版本的CCA点云图
    plot_2d_cca_separate(U_grouped_averaged, V_grouped_averaged, colors, args.nn_type)
    
    # 保存2D版本为SVG
    save_2d_cca_svg(U_grouped_averaged, V_grouped_averaged, colors, args.nn_type)

    # 绘制 Canonical modes (3D版本)
    for i in range(2):
        cm1, cm2, cm3 = i * 3, i * 3 + 1, i * 3 + 2

        fig = plt.figure()
        ax1 = fig.add_subplot(121, projection='3d')
        ax2 = fig.add_subplot(122, projection='3d')

        # 定义鼠标事件处理函数
        def on_mouse_event(event):
            if event.inaxes == ax1:
                ax2.view_init(elev=ax1.elev, azim=ax1.azim)
            elif event.inaxes == ax2:
                ax1.view_init(elev=ax2.elev, azim=ax2.azim)
            fig.canvas.draw()
        fig.canvas.mpl_connect('motion_notify_event', on_mouse_event)

        for j in range(10):
            ax1.scatter(U_grouped_averaged[j][:, cm1], U_grouped_averaged[j][:, cm2], U_grouped_averaged[j][:, cm3], c=colors[j], marker='o')
            ax2.scatter(V_grouped_averaged[j][:, cm1], V_grouped_averaged[j][:, cm2], V_grouped_averaged[j][:, cm3], c=colors[j], marker='o')

        ax1.set_title('Neural Dynamics')
        ax2.set_title('Solution Structure')

        ax1.set_xlabel(f'Canonical mode {cm1}')
        ax1.set_ylabel(f'Canonical mode {cm2}')
        ax1.set_zlabel(f'Canonical mode {cm3}')
        ax2.set_xlabel(f'Canonical mode {cm1}')
        ax2.set_ylabel(f'Canonical mode {cm2}')
        ax2.set_zlabel(f'Canonical mode {cm3}')

        # 设置两个 view 为 top-down 视角
        ax1.view_init(elev=90, azim=0)
        ax2.view_init(elev=90, azim=0)

        plt.show()
